import math

import cv2
import numpy as np


class Macrometer:
    def __init__(self, config):
        self.config = config

        self.stereoL = cv2.StereoSGBM_create(**config['SGBM']['left_param']) # 创建立体匹配的算法SGBM
        self.stereoR = cv2.StereoSGBM_create(**config['SGBM']['right_param'])

        self.param_camera_left = {}
        self.param_camera_right = {}

        self.param_translation_matrix_L2R = None
        self.param_rotate_matrix_L2R = None

        self.image_size = ()

        # 计算参数
        self.get_camera_LR_param()

    def get_camera_LR_param(self):
        # 单个相机的参数
        self.param_camera_left = self.calculate_camera_param(self.config['camera_left'])  #计算相机内参矩阵K 畸变系数distortion
        self.param_camera_right = self.calculate_camera_param(self.config['camera_right'])

        # 平移向量 外参矩阵：平移
        self.param_translation_matrix_L2R = \
            np.array(self.config['camera_right']['location'] ) - np.array(self.config['camera_left']['location'])

        # 旋转矩阵
        rotate_matrix_L = self.calculate_rotate_matrix(self.config['camera_left']['rotation'])
        rotate_matrix_R = self.calculate_rotate_matrix(self.config['camera_right']['rotation'])

        # 外参矩阵：旋转
        self.param_rotate_matrix_L2R = rotate_matrix_R @ np.transpose(rotate_matrix_L)  # @运算符： 矩阵乘法

        # 图像分辨率
        self.image_size = (
            self.config['camera_left']['resolution_x'],
            self.config['camera_left']['resolution_y'],
        )

    @staticmethod
    def calculate_camera_param(param_blender):   #计算相机内参矩阵K 畸变系数distortion
        param_camera = {}

        # 相机内参
        focal_x = param_blender['focal_length'] * param_blender['resolution_x'] / param_blender['sensor_width']
        focal_y = param_blender['focal_length'] * param_blender['resolution_y'] / param_blender['sensor_height']

        principal_point_x = param_blender['resolution_x'] / 2
        principal_point_y = param_blender['resolution_y'] / 2

        param_camera['K'] = np.array(
            [
                [focal_x, 0, principal_point_x],
                [0, focal_y, principal_point_y],
                [0, 0, 1],
            ])

        # 畸变系数
        param_camera['distortion'] = np.array(param_blender['distortion']).astype(np.float64)

        return param_camera

    @staticmethod
    def calculate_rotate_matrix(param):
        x, y, z = param

        R_x = np.array([
            [1, 0, 0],
            [0, np.cos(x), -np.sin(x)],
            [0, np.sin(x), np.cos(x)]
        ])

        R_y = np.array([
            [np.cos(y), 0, np.sin(y)],
            [0, 1, 0],
            [-np.sin(y), 0, np.cos(y)]
        ])

        R_z = np.array([
            [np.cos(z), -np.sin(z), 0],
            [np.sin(z), np.cos(z), 0],
            [0, 0, 1]
        ])

        # Combine the rotation matrices
        R = R_z @ R_y @ R_x

        return R

    def calculate_disparity(self, img_left, img_right):   # 计算视差图
        disparityLR = self.stereoL.compute(img_left, img_right).astype(np.float32) / 16.0   # 计算视差图
        disparityRL = self.stereoR.compute(img_right, img_left).astype(np.float32) / 16.0

        return disparityLR, disparityRL

    def calculate_3d_coordinate(self, disparity):
        # 要把视差图乘以掩码

        R1, R2, P1, P2, Q,  validPixROI1, validPixROI2 = cv2.stereoRectify(  # 得到矫正双目图像所需的变换矩阵和投影矩阵
            cameraMatrix1=self.param_camera_left['K'],
            distCoeffs1=self.param_camera_left['distortion'],
            cameraMatrix2=self.param_camera_right['K'],
            distCoeffs2=self.param_camera_right['distortion'],
            imageSize=self.image_size,
            R=self.param_rotate_matrix_L2R,
            T=self.param_translation_matrix_L2R
        )

        print(Q)  # 4x4的视差深度映射矩阵

        point_3d = cv2.reprojectImageTo3D(disparity=disparity, Q=Q, handleMissingValues=False)  # 视差图生成三维点云

        return point_3d

    @staticmethod
    def calculate_distance(point_clout, centroid):
        centroid_0, centroid_1 = centroid
        centroid_0 = (int(centroid_0[0]), int(centroid_0[1]))
        centroid_1 = (int(centroid_1[0]), int(centroid_1[1]))

        point_0 = np.array(point_clout[centroid_0[1]][centroid_0[0]])
        point_1 = np.array(point_clout[centroid_1[1]][centroid_1[0]])

        distance = np.linalg.norm(point_0 - point_1)

        return distance



